// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <melon/benchmark.h>
#include <melon/init/init.h>
#include <pollux/functions/macros.h>
#include <pollux/functions/registerer.h>
#include <pollux/functions/lib/benchmarks/FunctionBenchmarkBase.h>
#include <pollux/functions/prestosql/registration/registration_functions.h>
#include <pollux/vector/vector_fuzzer.h>

namespace {
using namespace kumo::pollux;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::functions;

template <typename TExec>
struct udf_bitwise_arithmetic_shift_right_nocheck {
  POLLUX_DEFINE_FUNCTION_TYPES(TExec);

  MELON_ALWAYS_INLINE bool
  call(int64_t& result, int64_t number, int64_t shift) {
    result = number >> shift;
    return true;
  }
};

template <typename TExec>
struct udf_bitwise_logical_shift_right_nocheck {
  POLLUX_DEFINE_FUNCTION_TYPES(TExec);

  MELON_ALWAYS_INLINE bool
  call(int64_t& result, int64_t number, int64_t shift, int64_t bits) {
    if (bits == 64) {
      result = number >> shift;
      return true;
    }

    result = (number & ((1LL << bits) - 1)) >> shift;
    return true;
  }
};

class BitwiseBenchmark : public functions::test::FunctionBenchmarkBase {
 public:
  BitwiseBenchmark() : FunctionBenchmarkBase() {
    functions::prestosql::registerBitwiseFunctions();
    register_function<
        udf_bitwise_arithmetic_shift_right_nocheck,
        int64_t,
        int64_t,
        int64_t>({"bitwise_arithmetic_shift_right_nocheck"});
    register_function<
        udf_bitwise_logical_shift_right_nocheck,
        int64_t,
        int64_t,
        int64_t,
        int64_t>({"bitwise_logical_shift_right_nocheck"});
  }

  void runBitwise(const std::string& fnName, bool logicalShift = false) {
    melon::BenchmarkSuspender suspender;

    VectorFuzzer::Options opts;
    opts.vectorSize = 100'000;
    VectorFuzzer fuzzer(opts, execCtx_.pool());
    auto vectorLeft = fuzzer.fuzzFlat(BIGINT());
    auto vectorRight = vectorMaker_.flat_vector<int32_t>(
        opts.vectorSize,
        [](auto row) {
          return row % 62 + 2; // Ensure value between [2, 64]
        },
        nullptr);

    auto row_vector = logicalShift
        ? vectorMaker_.row_vector({vectorLeft, vectorRight, vectorRight})
        : vectorMaker_.row_vector({vectorLeft, vectorRight});
    auto exprSet = logicalShift
        ? compileExpression(
              fmt::format("{}(c0, c1, c2)", fnName), row_vector->type())
        : compileExpression(
              fmt::format("{}(c0, c1)", fnName), row_vector->type());

    suspender.dismiss();

    doRun(exprSet, row_vector);
  }

  void doRun(ExprSet& exprSet, const RowVectorPtr& row_vector) {
    uint32_t cnt = 0;
    for (auto i = 0; i < 100; i++) {
      cnt += evaluate(exprSet, row_vector)->size();
    }
    melon::doNotOptimizeAway(cnt);
  }
};

BENCHMARK(bitwise_arithmetic_shift_right) {
  BitwiseBenchmark benchmark;
  benchmark.runBitwise("bitwise_arithmetic_shift_right");
}

BENCHMARK_RELATIVE(bitwise_arithmetic_shift_right_nocheck) {
  BitwiseBenchmark benchmark;
  benchmark.runBitwise("bitwise_arithmetic_shift_right_nocheck");
}

BENCHMARK_RELATIVE(bitwise_right_shift_arithmetic) {
  BitwiseBenchmark benchmark;
  benchmark.runBitwise("bitwise_right_shift_arithmetic");
}

BENCHMARK_RELATIVE(bitwise_right_shift) {
  BitwiseBenchmark benchmark;
  benchmark.runBitwise("bitwise_right_shift");
}

BENCHMARK_RELATIVE(bitwise_left_shift) {
  BitwiseBenchmark benchmark;
  benchmark.runBitwise("bitwise_left_shift");
}

BENCHMARK_RELATIVE(bitwise_xor) {
  BitwiseBenchmark benchmark;
  benchmark.runBitwise("bitwise_xor");
}

BENCHMARK_RELATIVE(bitwise_or) {
  BitwiseBenchmark benchmark;
  benchmark.runBitwise("bitwise_or");
}

BENCHMARK(bitwise_logical_shift_right) {
  BitwiseBenchmark benchmark;
  benchmark.runBitwise("bitwise_logical_shift_right", true);
}

BENCHMARK_RELATIVE(bitwise_logical_shift_right_nocheck) {
  BitwiseBenchmark benchmark;
  benchmark.runBitwise("bitwise_logical_shift_right_nocheck", true);
}

BENCHMARK_RELATIVE(bitwise_shift_left) {
  BitwiseBenchmark benchmark;
  benchmark.runBitwise("bitwise_shift_left", true);
}

} // namespace

int main(int argc, char** argv) {
  melon::Init init{&argc, &argv};
  memory::MemoryManager::initialize({});

  melon::runBenchmarks();
  return 0;
}
